#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
智能向量搜索工具
功能：输入文字，自动在Docker Milvus向量数据库中搜索相似的问题并返回结果
"""

import numpy as np
from sklearn.feature_extraction.text import TfidfVectorizer
import jieba
import re
from pymilvus import connections, Collection
import sys


class SmartVectorSearch:
    def __init__(self, collection_name="knowledge_lines", host="localhost", port="19530"):
        self.collection_name = collection_name
        self.host = host
        self.port = port
        self.collection = None
        self.vectorizer = None
        self.is_connected = False
        
    def connect_to_milvus(self):
        """连接到Milvus数据库"""
        try:
            connections.connect("default", host=self.host, port=self.port)
            self.collection = Collection(self.collection_name)
            self.collection.load()
            print(f"✅ 成功连接到Milvus服务器 {self.host}:{self.port}")
            print(f"✅ 加载集合: {self.collection_name} (共{self.collection.num_entities}条记录)")
            self.is_connected = True
            return True
        except Exception as e:
            print(f"❌ 连接失败: {e}")
            return False
    
    def preprocess_text(self, text):
        """文本预处理"""
        # 移除特殊字符和数字
        text = re.sub(r'[^\w\s]', ' ', text)
        text = re.sub(r'\d+', ' ', text)
        text = text.lower()
        
        # 中文分词
        words = jieba.cut(text)
        text = ' '.join(words)
        
        # 移除多余空格
        text = re.sub(r'\s+', ' ', text).strip()
        return text
    
    def build_vectorizer(self, max_features=100):
        """构建向量化器（需要用训练数据拟合）"""
        try:
            # 获取一些样本数据来训练向量化器
            sample_results = self.collection.query(
                expr='line_number > 0',
                limit=1000,  # 使用1000条样本训练
                output_fields=['line_text']
            )
            
            if not sample_results:
                print("❌ 无法获取样本数据来训练向量化器")
                return False
            
            # 预处理样本文本
            sample_texts = [self.preprocess_text(result['line_text']) for result in sample_results]
            sample_texts = [text for text in sample_texts if text.strip()]  # 过滤空文本
            
            if not sample_texts:
                print("❌ 没有有效的样本文本")
                return False
            
            # 训练TF-IDF向量化器
            self.vectorizer = TfidfVectorizer(max_features=max_features, ngram_range=(1, 2))
            self.vectorizer.fit(sample_texts)
            
            print(f"✅ 向量化器训练完成 (特征维度: {len(self.vectorizer.get_feature_names_out())})")
            return True
            
        except Exception as e:
            print(f"❌ 构建向量化器失败: {e}")
            return False
    
    def vectorize_query(self, query_text):
        """将查询文本向量化"""
        if not self.vectorizer:
            print("❌ 向量化器未初始化")
            return None
        
        try:
            processed_query = self.preprocess_text(query_text)
            if not processed_query.strip():
                print("❌ 查询文本预处理后为空")
                return None
            
            query_vector = self.vectorizer.transform([processed_query]).toarray()
            return query_vector[0]
            
        except Exception as e:
            print(f"❌ 查询向量化失败: {e}")
            return None
    
    def search_similar(self, query_text, top_k=5, score_threshold=None):
        """搜索相似文本"""
        if not self.is_connected:
            print("❌ 未连接到Milvus")
            return []
        
        try:
            # 向量化查询文本
            query_vector = self.vectorize_query(query_text)
            if query_vector is None:
                return []
            
            # 执行向量搜索
            search_params = {"metric_type": "L2", "params": {"nprobe": 16}}
            
            results = self.collection.search(
                [query_vector.tolist()],
                "vector",
                search_params,
                limit=top_k,
                output_fields=["filename", "line_number", "line_text"]
            )
            
            # 处理搜索结果
            similar_results = []
            for hits in results:
                for hit in hits:
                    # 计算相似度分数 (距离越小越相似)
                    distance = hit.distance
                    similarity_score = 1 / (1 + distance)  # 转换为相似度分数
                    
                    # 应用分数阈值过滤
                    if score_threshold is None or similarity_score >= score_threshold:
                        similar_results.append({
                            'filename': hit.entity.get('filename'),
                            'line_number': hit.entity.get('line_number'),
                            'line_text': hit.entity.get('line_text'),
                            'distance': distance,
                            'similarity_score': similarity_score
                        })
            
            return similar_results
            
        except Exception as e:
            print(f"❌ 搜索失败: {e}")
            return []
    
    def keyword_search(self, query_text, limit=10):
        """关键词搜索作为补充"""
        try:
            # 提取关键词
            keywords = query_text.split()
            if not keywords:
                return []
            
            # 构建搜索表达式
            conditions = []
            for keyword in keywords[:3]:  # 最多使用前3个关键词
                if len(keyword) > 1:  # 过滤单字符
                    conditions.append(f'line_text like "%{keyword}%"')
            
            if not conditions:
                return []
            
            expr = ' or '.join(conditions)
            
            results = self.collection.query(
                expr=expr,
                limit=limit,
                output_fields=['filename', 'line_number', 'line_text']
            )
            
            return results
            
        except Exception as e:
            print(f"❌ 关键词搜索失败: {e}")
            return []
    
    def smart_search(self, query_text, top_k=5, include_keyword_search=True):
        """智能搜索：结合向量搜索和关键词搜索"""
        print(f"\n🔍 搜索查询: '{query_text}'")
        print("=" * 60)
        
        all_results = []
        
        # 1. 向量相似性搜索
        print("📊 执行向量相似性搜索...")
        vector_results = self.search_similar(query_text, top_k)
        
        if vector_results:
            print(f"✅ 找到 {len(vector_results)} 条相似结果")
            all_results.extend([{**r, 'search_type': 'vector'} for r in vector_results])
        else:
            print("⚠️  向量搜索无结果")
        
        # 2. 关键词搜索（作为补充）
        if include_keyword_search:
            print("\n🔤 执行关键词搜索...")
            keyword_results = self.keyword_search(query_text, limit=5)
            
            if keyword_results:
                print(f"✅ 找到 {len(keyword_results)} 条关键词匹配结果")
                # 转换格式并添加到结果中
                for result in keyword_results:
                    all_results.append({
                        'filename': result['filename'],
                        'line_number': result['line_number'],
                        'line_text': result['line_text'],
                        'search_type': 'keyword',
                        'similarity_score': 0.5  # 给关键词搜索一个默认分数
                    })
            else:
                print("⚠️  关键词搜索无结果")
        
        # 3. 去重和排序
        seen = set()
        unique_results = []
        for result in all_results:
            key = (result['filename'], result['line_number'])
            if key not in seen:
                seen.add(key)
                unique_results.append(result)
        
        # 按相似度分数排序
        unique_results.sort(key=lambda x: x.get('similarity_score', 0), reverse=True)
        
        return unique_results[:top_k * 2]  # 返回更多结果
    
    def display_results(self, results, query_text):
        """显示搜索结果"""
        if not results:
            print("\n❌ 没有找到相关结果")
            return
        
        print(f"\n🎯 搜索结果 (共{len(results)}条):")
        print("=" * 60)
        
        for i, result in enumerate(results, 1):
            print(f"\n--- 结果 {i} ---")
            print(f"📁 文件: {result['filename']}")
            print(f"📍 行号: {result['line_number']}")
            print(f"🔍 搜索类型: {result['search_type']}")
            
            if 'similarity_score' in result:
                print(f"📊 相似度: {result['similarity_score']:.3f}")
            
            # 高亮显示匹配的关键词
            text = result['line_text']
            for keyword in query_text.split():
                if len(keyword) > 1:
                    text = text.replace(keyword, f"【{keyword}】")
            
            print(f"📝 内容: {text}")
            
            if i >= 10:  # 最多显示10条
                print(f"\n... 还有 {len(results) - 10} 条结果未显示")
                break


def interactive_search():
    """交互式搜索"""
    print("🤖 智能向量搜索工具")
    print("=" * 50)
    print("功能：输入问题，自动搜索知识库中的相似内容")
    print("=" * 50)
    
    # 初始化搜索器
    searcher = SmartVectorSearch()
    
    # 连接数据库
    if not searcher.connect_to_milvus():
        print("❌ 无法连接到Milvus，程序退出")
        return
    
    # 构建向量化器
    print("\n🔧 正在初始化向量化器...")
    if not searcher.build_vectorizer():
        print("❌ 向量化器初始化失败，程序退出")
        return
    
    print("\n✅ 系统准备就绪！")
    print("💡 提示：输入 'quit' 或 'exit' 退出程序")
    
    # 开始交互式搜索
    while True:
        print("\n" + "=" * 50)
        query = input("🔍 请输入您的问题: ").strip()
        
        if query.lower() in ['quit', 'exit', '退出']:
            print("👋 再见！")
            break
        
        if not query:
            print("⚠️  请输入有效的问题")
            continue
        
        try:
            # 执行智能搜索
            results = searcher.smart_search(query, top_k=5)
            
            # 显示结果
            searcher.display_results(results, query)
            
        except KeyboardInterrupt:
            print("\n\n👋 用户中断，再见！")
            break
        except Exception as e:
            print(f"\n❌ 搜索过程中出现错误: {e}")


def command_line_search():
    """命令行搜索"""
    if len(sys.argv) < 2:
        print("用法: python smart_search.py '您的问题'")
        print("或者: python smart_search.py  # 进入交互模式")
        return
    
    query = ' '.join(sys.argv[1:])
    
    # 初始化搜索器
    searcher = SmartVectorSearch()
    
    # 连接和初始化
    if not searcher.connect_to_milvus():
        return
    
    if not searcher.build_vectorizer():
        return
    
    # 执行搜索
    results = searcher.smart_search(query, top_k=5)
    searcher.display_results(results, query)


if __name__ == "__main__":
    if len(sys.argv) > 1:
        command_line_search()
    else:
        interactive_search()
